#ifndef GEDGEE_ARM_H
#define GEDGEE_ARM_H

#include "../onlycpp.h"
#include <arm_neon.h>

namespace gedgee {

// NEON类型转换工具
inline uint8x16_t to_uint8x16(vector128 v) {
    return vloadq_u8(v.u8);
}

inline uint16x8_t to_uint16x8(vector128 v) {
    return vloadq_u16(v.u16);
}

inline uint32x4_t to_uint32x4(vector128 v) {
    return vloadq_u32(v.u32);
}

inline uint64x2_t to_uint64x2(vector128 v) {
    return vloadq_u64(v.u64);
}

inline int8x16_t to_int8x16(vector128 v) {
    return vloadq_s8(v.i8);
}

inline int16x8_t to_int16x8(vector128 v) {
    return vloadq_s16(v.i16);
}

inline int32x4_t to_int32x4(vector128 v) {
    return vloadq_s32(v.i32);
}

inline int64x2_t to_int64x2(vector128 v) {
    return vloadq_s64(v.i64);
}

inline float32x4_t to_float32x4(vector128 v) {
    return vld1q_f32(v.f32);
}

inline float64x2_t to_float64x2(vector128 v) {
    return vld1q_f64(v.f64);
}

template <typename T>
inline vector128 to_vector128(T neon) {
    vector128 v;
    vstore(neon, 0, reinterpret_cast<typename std::remove_reference<decltype(neon)>::type::element_type*>(&v));
    return v;
}

// 类型转换特化
template <>
inline vector128 v128_cvt<u8, u16>(vector128 vec) {
    uint8x16_t input = to_uint8x16(vec);
    return to_vector128(vmovl_u8(vget_low_u8(input)));
}

template <>
inline vector128 v128_cvt<u16, u8>(vector128 vec) {
    uint16x8_t input = to_uint16x8(vec);
    uint8x8_t low = vqmovn_u16(input);
    return to_vector128(vcombine_u8(low, low));
}

template <>
inline vector128 v128_cvt<i32, f32>(vector128 vec) {
    int32x4_t input = to_int32x4(vec);
    return to_vector128(vcvtq_f32_s32(input));
}

template <>
inline vector128 v128_cvt<f32, i32>(vector128 vec) {
    float32x4_t input = to_float32x4(vec);
    return to_vector128(vcvtq_s32_f32(input));
}

// 存储函数特化
template <>
inline void v128_store<u8>(u8 *p, vector128 vec) {
    vst1q_u8(p, to_uint8x16(vec));
}

template <>
inline void v128_store<u16>(u16 *p, vector128 vec) {
    vst1q_u16(p, to_uint16x8(vec));
}

template <>
inline void v128_store<u32>(u32 *p, vector128 vec) {
    vst1q_u32(p, to_uint32x4(vec));
}

template <>
inline void v128_store<u64>(u64 *p, vector128 vec) {
    vst1q_u64(p, to_uint64x2(vec));
}

template <>
inline void v128_store<i8>(i8 *p, vector128 vec) {
    vst1q_s8(p, to_int8x16(vec));
}

template <>
inline void v128_store<i16>(i16 *p, vector128 vec) {
    vst1q_s16(p, to_int16x8(vec));
}

template <>
inline void v128_store<i32>(i32 *p, vector128 vec) {
    vst1q_s32(p, to_int32x4(vec));
}

template <>
inline void v128_store<i64>(i64 *p, vector128 vec) {
    vst1q_s64(p, to_int64x2(vec));
}

template <>
inline void v128_store<f32>(f32 *p, vector128 vec) {
    vst1q_f32(p, to_float32x4(vec));
}

template <>
inline void v128_store<f64>(f64 *p, vector128 vec) {
    vst1q_f64(p, to_float64x2(vec));
}

// 设置函数特化
template <>
inline vector128 v128_set<u8>(const u8 n) {
    return to_vector128(vdupq_n_u8(n));
}

template <>
inline vector128 v128_set<u16>(const u16 n) {
    return to_vector128(vdupq_n_u16(n));
}

template <>
inline vector128 v128_set<u32>(const u32 n) {
    return to_vector128(vdupq_n_u32(n));
}

template <>
inline vector128 v128_set<u64>(const u64 n) {
    return to_vector128(vdupq_n_u64(n));
}

template <>
inline vector128 v128_set<i8>(const i8 n) {
    return to_vector128(vdupq_n_s8(n));
}

template <>
inline vector128 v128_set<i16>(const i16 n) {
    return to_vector128(vdupq_n_s16(n));
}

template <>
inline vector128 v128_set<i32>(const i32 n) {
    return to_vector128(vdupq_n_s32(n));
}

template <>
inline vector128 v128_set<i64>(const i64 n) {
    return to_vector128(vdupq_n_s64(n));
}

template <>
inline vector128 v128_set<f32>(const f32 n) {
    return to_vector128(vdupq_n_f32(n));
}

template <>
inline vector128 v128_set<f64>(const f64 n) {
    return to_vector128(vdupq_n_f64(n));
}

// 算术运算特化
template <>
inline vector128 v128_add<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_add<u16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_u16(to_uint16x8(vec1), to_uint16x8(vec2)));
}

template <>
inline vector128 v128_add<u32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_u32(to_uint32x4(vec1), to_uint32x4(vec2)));
}

template <>
inline vector128 v128_add<u64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_u64(to_uint64x2(vec1), to_uint64x2(vec2)));
}

template <>
inline vector128 v128_add<i8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_s8(to_int8x16(vec1), to_int8x16(vec2)));
}

template <>
inline vector128 v128_add<i16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_s16(to_int16x8(vec1), to_int16x8(vec2)));
}

template <>
inline vector128 v128_add<i32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_s32(to_int32x4(vec1), to_int32x4(vec2)));
}

template <>
inline vector128 v128_add<i64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_s64(to_int64x2(vec1), to_int64x2(vec2)));
}

template <>
inline vector128 v128_add<f32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_f32(to_float32x4(vec1), to_float32x4(vec2)));
}

template <>
inline vector128 v128_add<f64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vaddq_f64(to_float64x2(vec1), to_float64x2(vec2)));
}

// 减法运算特化
template <>
inline vector128 v128_sub<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_sub<u16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_u16(to_uint16x8(vec1), to_uint16x8(vec2)));
}

template <>
inline vector128 v128_sub<u32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_u32(to_uint32x4(vec1), to_uint32x4(vec2)));
}

template <>
inline vector128 v128_sub<u64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_u64(to_uint64x2(vec1), to_uint64x2(vec2)));
}

template <>
inline vector128 v128_sub<i8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_s8(to_int8x16(vec1), to_int8x16(vec2)));
}

template <>
inline vector128 v128_sub<i16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_s16(to_int16x8(vec1), to_int16x8(vec2)));
}

template <>
inline vector128 v128_sub<i32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_s32(to_int32x4(vec1), to_int32x4(vec2)));
}

template <>
inline vector128 v128_sub<i64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_s64(to_int64x2(vec1), to_int64x2(vec2)));
}

template <>
inline vector128 v128_sub<f32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_f32(to_float32x4(vec1), to_float32x4(vec2)));
}

template <>
inline vector128 v128_sub<f64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vsubq_f64(to_float64x2(vec1), to_float64x2(vec2)));
}

// 乘法运算特化
template <>
inline vector128 v128_mul<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_mul<u16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_u16(to_uint16x8(vec1), to_uint16x8(vec2)));
}

template <>
inline vector128 v128_mul<u32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_u32(to_uint32x4(vec1), to_uint32x4(vec2)));
}

template <>
inline vector128 v128_mul<i8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_s8(to_int8x16(vec1), to_int8x16(vec2)));
}

template <>
inline vector128 v128_mul<i16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_s16(to_int16x8(vec1), to_int16x8(vec2)));
}

template <>
inline vector128 v128_mul<i32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_s32(to_int32x4(vec1), to_int32x4(vec2)));
}

template <>
inline vector128 v128_mul<f32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_f32(to_float32x4(vec1), to_float32x4(vec2)));
}

template <>
inline vector128 v128_mul<f64>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmulq_f64(to_float64x2(vec1), to_float64x2(vec2)));
}

// 最大值/最小值运算特化
template <>
inline vector128 v128_max<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_max<u16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_u16(to_uint16x8(vec1), to_uint16x8(vec2)));
}

template <>
inline vector128 v128_max<u32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_u32(to_uint32x4(vec1), to_uint32x4(vec2)));
}

template <>
inline vector128 v128_max<i8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_s8(to_int8x16(vec1), to_int8x16(vec2)));
}

template <>
inline vector128 v128_max<i16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_s16(to_int16x8(vec1), to_int16x8(vec2)));
}

template <>
inline vector128 v128_max<i32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_s32(to_int32x4(vec1), to_int32x4(vec2)));
}

template <>
inline vector128 v128_max<f32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vmaxq_f32(to_float32x4(vec1), to_float32x4(vec2)));
}

template <>
inline vector128 v128_min<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_min<u16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_u16(to_uint16x8(vec1), to_uint16x8(vec2)));
}

template <>
inline vector128 v128_min<u32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_u32(to_uint32x4(vec1), to_uint32x4(vec2)));
}

template <>
inline vector128 v128_min<i8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_s8(to_int8x16(vec1), to_int8x16(vec2)));
}

template <>
inline vector128 v128_min<i16>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_s16(to_int16x8(vec1), to_int16x8(vec2)));
}

template <>
inline vector128 v128_min<i32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_s32(to_int32x4(vec1), to_int32x4(vec2)));
}

template <>
inline vector128 v128_min<f32>(vector128 vec1, vector128 vec2) {
    return to_vector128(vminq_f32(to_float32x4(vec1), to_float32x4(vec2)));
}

// 位运算特化
template <>
inline vector128 v128_and<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vandq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_or<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vorrq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_xor<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(veorq_u8(to_uint8x16(vec1), to_uint8x16(vec2)));
}

template <>
inline vector128 v128_andnot<u8>(vector128 vec1, vector128 vec2) {
    return to_vector128(vbicq_u8(to_uint8x16(vec2), to_uint8x16(vec1)));
}

// 移位运算特化
template <>
inline vector128 v128_sll<u16>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_u16(to_uint16x8(vec), count));
}

template <>
inline vector128 v128_sll<u32>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_u32(to_uint32x4(vec), count));
}

template <>
inline vector128 v128_sll<u64>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_u64(to_uint64x2(vec), count));
}

template <>
inline vector128 v128_sll<i16>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_s16(to_int16x8(vec), count));
}

template <>
inline vector128 v128_sll<i32>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_s32(to_int32x4(vec), count));
}

template <>
inline vector128 v128_sll<i64>(vector128 vec, uint32_t count) {
    return to_vector128(vshlq_n_s64(to_int64x2(vec), count));
}

template <>
inline vector128 v128_srl<u16>(vector128 vec, uint32_t count) {
    return to_vector128(vshrq_n_u16(to_uint16x8(vec), count));
}

template <>
inline vector128 v128_srl<u32>(vector128 vec, uint32_t count) {
    return to_vector128(vshrq_n_u32(to_uint32x4(vec), count));
}

template <>
inline vector128 v128_srl<u64>(vector128 vec, uint32_t count) {
    return to_vector128(vshrq_n_u64(to_uint64x2(vec), count));
}

template <>
inline vector128 v128_sra<i16>(vector128 vec, uint32_t count) {
    return to_vector128(vshrq_n_s16(to_int16x8(vec), count));
}

template <>
inline vector128 v128_sra<i32>(vector128 vec, uint32_t count) {
    return to_vector128(vshrq_n_s32(to_int32x4(vec), count));
}

// 数学函数特化
template <>
inline vector128 v128_abs<i8>(vector128 vec) {
    return to_vector128(vabsq_s8(to_int8x16(vec)));
}

template <>
inline vector128 v128_abs<i16>(vector128 vec) {
    return to_vector128(vabsq_s16(to_int16x8(vec)));
}

template <>
inline vector128 v128_abs<i32>(vector128 vec) {
    return to_vector128(vabsq_s32(to_int32x4(vec)));
}

template <>
inline vector128 v128_abs<f32>(vector128 vec) {
    return to_vector128(vabsq_f32(to_float32x4(vec)));
}

template <>
inline vector128 v128_neg<i8>(vector128 vec) {
    return to_vector128(vnegq_s8(to_int8x16(vec)));
}

template <>
inline vector128 v128_neg<i16>(vector128 vec) {
    return to_vector128(vnegq_s16(to_int16x8(vec)));
}

template <>
inline vector128 v128_neg<i32>(vector128 vec) {
    return to_vector128(vnegq_s32(to_int32x4(vec)));
}

template <>
inline vector128 v128_neg<f32>(vector128 vec) {
    return to_vector128(vnegq_f32(to_float32x4(vec)));
}

template <>
inline vector128 v128_sqrt<f32>(vector128 vec) {
    return to_vector128(vsqrtq_f32(to_float32x4(vec)));
}

template <>
inline vector128 v128_rsqrt<f32>(vector128 vec) {
    float32x4_t estimate = vrsqrteq_f32(to_float32x4(vec));
    return to_vector128(vmulq_f32(estimate, vrsqrtsq_f32(vec, vmulq_f32(estimate, estimate))));
}

template <>
inline vector128 v128_rcp<f32>(vector128 vec) {
    float32x4_t estimate = vrecpeq_f32(to_float32x4(vec));
    return to_vector128(vmulq_f32(estimate, vrecpsq_f32(estimate, vec)));
}

// 舍入函数特化
template <>
inline vector128 v128_round<f32>(vector128 vec) {
    return to_vector128(vrndnq_f32(to_float32x4(vec)));
}

} // namespace gedgee

#endif // GEDGEE_ARM_H